# 消费者代码，测试可以使用：python consume.py error
import functools
import threading
import time
# from pathos.multiprocessing import Pool
from multiprocessing import Pool, Manager
from typing import Optional, List, Union, Callable

from loguru import logger
import pika
import sys

Flag_dict = Manager().dict()


class Run_MUlPro:
    def __init__(self, poolnums):
        self.poolnums = poolnums
        self.create_pool()

    def create_pool(self):
        self.pool = Pool(self.poolnums)

    def process(self, method, args):
        self.pool.apply_async(method, args=args)
        # method(*args)

    def __del__(self):
        self.pool.close()
        self.pool.join()


class Consumer:
    def __init__(self, host: str, port: int = 5672, virtual_host: str = '/', username: str = 'guest',
                 password: str = 'guest', consumernums=1, *args, **kwargs):
        self.consumer_nums = consumernums
        self.connection = self._connect_mq(
            host, port, virtual_host, username, password)
        if self.connection:
            self.channel = self.connection.channel()
        self.run_pro = Run_MUlPro(consumernums +2)
        self.exchange = kwargs['exchange'] if 'exchange' in kwargs else 'default_topic_messages'
        self.ack_objs = {}
        exchange_durable = kwargs['exchange_durable'] if 'exchange_durable' in kwargs else False
        self.channel.exchange_declare(
            exchange=self.exchange,
            exchange_type='topic',
            durable=exchange_durable
        )

    def _connect_mq(self, host: str, port: int = 5672, virtual_host: str = '/',
                    username: str = 'guest', password: str = 'guest') -> Optional[pika.BlockingConnection]:
        """connect to rabbitMQ.

        :params str host: host for rabbitMQ.
        :params int post: port for rabbitMQ, default is 5672.
        :params str virtual_host: virtual_host for rabbitMQ, default is '/'.
        :params str username: username for rabbit credential, default is 'guest'.
        :params str password: password for rabbit credential, default is 'guest'.

        :return: a BlockingConnection object or None
        """
        try:
            credentials = pika.PlainCredentials(username, password)
            parameters = pika.ConnectionParameters(
                host, port, virtual_host, credentials)
            connection = pika.BlockingConnection(parameters)
            return connection
        except Exception as e:
            logger.error(f'RabbitMQ Connect failed: {e}')
            return None

    def __del__(self):
        self.connection.close()


    def ack_message(self, sleep_time=0.01):
        """Note that `ch` must be the same pika channel instance via which
        the message being ACKed was retrieved (AMQP protocol constraint).
        """
        while True:
            lists = []
            for k, v in Flag_dict.items():
                if v == 2:
                    cb = functools.partial(self.partial_ack_method, self.ch, k)
                    self.connection.add_callback_threadsafe(cb)
                    # self.ch.basic_ack(k)
                    logger.info('%s 发送ack确认' % k)
                    lists.append(k)
            for i in lists:
                if Flag_dict:
                    Flag_dict.pop(i)
            time.sleep(sleep_time)


    def partial_ack_method(self, ch, delivery_tag):
        """Note that `ch` must be the same pika channel instance via which
        the message being ACKed was retrieved (AMQP protocol constraint).
        """
        if ch.is_open:
            ch.basic_ack(delivery_tag)
        else:
            # Channel is already closed, so we can't ACK this message;
            # log and/or do something that makes sense for your app in this case.
            pass

    @staticmethod
    def add_callback(consumer_key, callback, body):
        # connection, callback, body,
        Flag_dict[consumer_key] = 1
        callback(body)
        Flag_dict[consumer_key] = 2


    def on_message(self, ch, method_frame, _header_frame, body, args):
        callback = args[0]
        # delivery_tag = method_frame.delivery_tag
        # cb = functools.partial(self.ack_message1, ch, delivery_tag)
        consumer_key = method_frame.delivery_tag
        Flag_dict[consumer_key] = 0
        self.ch = ch

        self.run_pro.process(Consumer.add_callback, args=(consumer_key, callback, body))


    def process_message(
            self, queue: str, binding_keys: Union[str, List[str]], ttl_minutes: int = 360,  callback: Callable = None) -> None:
        """Consumers start to consume, using this method will start to get content from the message queue and process it.

        :params str queue: Queue name.
        :params str|list binding_keys: A list containing multiple binding keys.
            If only one key is bound, a string can also be passed in.
        :params int ttl_second: Message expiration time.
        :params callable callback: The callback function that needs to be executed.
            Default method will only print the key and body of the message, you need to write this function by yourself.
            The format of the callback function is like this:

                def callback(message_body):
                    pass
        """
        ttl = ttl_minutes * 1000 * 60
        result = self.channel.queue_declare(
            queue, durable=True, arguments={'x-message-ttl': ttl})
        queue_name = result.method.queue

        # convert str to list
        if isinstance(binding_keys, str):
            binding_keys = [binding_keys]

        for binding_key in binding_keys:
            self.channel.queue_bind(
                exchange=self.exchange, queue=queue_name, routing_key=binding_key
            )

        logger.info('Waiting for task. To exit press CTRL+C')
        self.channel.basic_qos(prefetch_count=self.consumer_nums)
        on_message_callback = functools.partial(
            self.on_message, args=(callback, )
        )
        self.channel.basic_consume(
            queue=queue_name,
            on_message_callback=on_message_callback
        )
        # start consuming
        ts = threading.Thread(target=self.ack_message)
        ts.start()
        try:
            self.channel.start_consuming()
        except KeyboardInterrupt:
            ts.join()
            self.channel.stop_consuming()


def demo(body):
    logger.debug(body)
    time.sleep(3)
    logger.debug(str(body) + '<<<')


class C:
    def pro(self):
        self.r = Run_MUlPro(2)
        self.r.create_pool()
        self.r.process(demo, args=('sdf', ))

if __name__ == '__main__':
    consumer = Consumer(host='152.136.187.200', username='admin', password='admin', exchange='topic_logs')
    consumer.process_message('ce_shi', binding_keys='ce.shi', callback=demo)
    # r = Run_MUlPro(4)
    # r.process(add_ack, args=('sdf','sdf','sdf', C().pro))
